import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import urlparse, parse_qs
import json
import threading
import sys
from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import FormattedText
from rich import print
from rich.panel import Panel
from rich.text import Text
token_store = None

def login_via_web_app(greeting=True):
    try:
        if greeting:
            message = Text.assemble(
                ("\n🛡️ Welcome to Pentest Muse!\n\n", "bold magenta"),
                ("Your personal assistant for penetration testing insights and guidance.\n\n", "bold cyan"),
                ("Let's get you set up:\n\n", "bold yellow"),
                ("🌐 Web App Login (recommended):\n", "bold green"),
                ("To enjoy a seamless experience with extended features, please authenticate through our web app. Just navigate to: ", "bold"),
                ("https://www.pentestmuse.ai/cli/login\n\n", "bold blue"),
                ("🔑 OpenAI API Key:\n", "bold red"),
                ("If you prefer to use your own OpenAI API keys, please restart with argument --openai-api-key=[your openai api key].\n", "bold")
            )
            print(Panel(message))
        else:
            print_formatted_text(FormattedText([("bold", f"\nTo login, please open the following URL in a browswer window: "), ("bold blue", f"https://www.pentestmuse.ai/cli/login\n")]))

        # Start a local server to listen for the redirect
        server_address = ('localhost', 8024)
        shutdown_event = threading.Event()
        httpd = HTTPServer(server_address, RequestHandler)
        httpd.shutdown_event = shutdown_event
        httpd.token = None

        # Start the server in a new thread
        server_thread = threading.Thread(target=httpd.serve_forever)
        server_thread.daemon = True
        server_thread.start()

        # Wait for the server to shut down
        shutdown_event.wait()

        # Indicate that login was successful
        print_formatted_text(FormattedText([("bold", f"Login successful!\n")]))

        # Return the token
        return token_store
    except KeyboardInterrupt:
        print('Keyboard interrupt received, exiting.')
        sys.exit()


class RequestHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        # Do nothing to suppress all log messages
        pass

    def do_OPTIONS(self):
        self.send_response(200, "ok")
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header("Access-Control-Allow-Headers", "X-Requested-With, Content-type")
        self.end_headers()

    def do_POST(self):
        # Get the size of data
        content_length = int(self.headers['Content-Length'])

        # Get the data itself
        post_data = self.rfile.read(content_length)

        # Parse the JSON data
        data = json.loads(post_data)

        # Extract the token and id
        token = data.get('token')
        id = data.get('id')

        if token is not None and id is not None:
            global token_store
            token_store = token
        
            # Signal that the server has been shut down
            self.server.shutdown_event.set()

            # Stop the server
            self.server.shutdown()

        # Send a response to the web app
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b'OK')

